'''
Description: 
version: 
Author: aps.auto
Date: 2025-03-25 10:26:02
LastEditors: Luoxl
LastEditTime: 2025-03-25 16:10:17
'''
from airbot_py import airbot_play, airbot_grpc
from utils_py.utils_positions import *
import utils_py.utils_gamepad_limx as limx_gamepad
from functools import partial
import time
import json
import threading
import limxsdk.datatypes as datatypes

# lock = threading.Lock()
is_recording = False
# left_arm = airbot_play.create_agent(ip="localhost", port=50000)
right_arm = airbot_play.create_agent(ip="localhost", port=50001)
robot = limx_gamepad.initialize_robot("10.192.1.2")

class JoyReceiver(limx_gamepad.LimxReceiver):
    def __init__(self):
        super().__init__()

    def play_action(self, action):
        super().play_action(action)
        
        if action is None or action == "None":
            print("Action is not configured. Please try again.")
            return 

        elif action == "record":
            is_recording = True
            record(right_arm, "temp_action")
            return
        elif action == "stop":
            is_recording = False
            return
        elif action == "replay":
            play_action(right_arm, "temp_action")
            return
    def sensorJoyCallback(self, sensor_joy: datatypes.SensorJoy, config_path='record_config.ini'):
        super().sensorJoyCallback(sensor_joy, config_path)
        
# Create an instance of RobotReceiver to handle callbacks
receiver = JoyReceiver()
# Create partial functions for callbacks
sensorJoyCallback = partial(receiver.sensorJoyCallback)
robot.subscribeSensorJoy(sensorJoyCallback)
# base_dir = "recorded_actions"

def record(robot:airbot_play.AirbotPlay, actionName:str):
    joint_pose = []
    
    while True:
        joint_pose.append(robot.get_current_joint_q())
        time.sleep(0.01)
        if not is_recording:
            break

    with open("recorded_actions/" + actionName+".json", "w") as f:    
        json.dump(joint_pose, f)
    print("Recording completed.")

def play_action(robot:airbot_play.AirbotPlay, actionName:str):
    with open("recorded_actions/" + actionName+".json", "r") as f:
        joint_pose = json.load(f)
    for pose in joint_pose:
        robot.set_target_joint_q_waypoints(pose)
        time.sleep(0.01)

